#/**
 * @file main.cpp
 *
 * Copyright (C) 2024. Huawei Technologies Co., Ltd. All rights reserved.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 */
#include "scatter_reduce_tiling.h"
#include "data_utils.h"
#ifndef ASCENDC_CPU_DEBUG
#include "acl/acl.h"
#include "aclrtlaunch_scatter_reduce.h"
#else
#include "tikicpulib.h"
extern "C" __global__ __aicore__ void scatter_reduce(GM_ADDR self, GM_ADDR index, GM_ADDR src, GM_ADDR currentWorkspace, ScatterReduceTilingData tiling);
#endif

int32_t main(int32_t argc, char *argv[])
{
    uint32_t size_of_data_type = sizeof(DTYPE_SELF);
    const uint32_t blockNum = 8;
    const std::vector<uint32_t> selfShape = {9, 2};
    const std::vector<uint32_t> indexShape = {9, 12};
    const std::string reduce_str = "mean";
    const bool include_self = true;
    const uint32_t dim = 1;
    // const uint32_t BLOCK_SIZE = 64;

    // Total length of input tensors
    uint32_t selfTotalLength = 1;
    for (auto shape_dim : selfShape) {
        selfTotalLength *= shape_dim;
    }
    uint32_t indexTotalLength = 1;
    for (auto shape_dim : indexShape) {
        indexTotalLength *= shape_dim;
    }
    
    // Convert reduce string to enum value
    ReduceType reduce;
    if (reduce_str == "sum" || reduce_str == "SUM") {
        reduce = SUM;
    } else if (reduce_str == "prod" || reduce_str == "PROD") {
        reduce = PROD;
    } else if (reduce_str == "mean" || reduce_str == "MEAN") {
        reduce = MEAN;
    } else if (reduce_str == "amax" || reduce_str == "AMAX") {
        reduce = AMAX;
    } else if (reduce_str == "amin" || reduce_str == "AMIN") {
        reduce = AMIN;
    } else {
        // Default to sum
        reduce = SUM;
    }

    // Calculate tiling parameters
    uint32_t prevDimSize = 1;
    uint32_t selfDimSize = selfShape[dim];
    uint32_t indexDimSize = indexShape[dim];
    uint32_t postDimSize = 1;
    for (uint32_t i = 0; i < dim; i++) {
        prevDimSize *= selfShape[i];
    }
    for (uint32_t i = dim + 1; i < selfShape.size(); i++) {
        postDimSize *= selfShape[i];
    }
    uint32_t blockLength = (prevDimSize + blockNum - 1) / blockNum;
    uint32_t effectiveBlockNum = (prevDimSize + blockLength - 1) / blockLength;
    uint32_t lastBlockLength = prevDimSize - (effectiveBlockNum - 1) * blockLength;

    // Calculate buffer sizes
    size_t selfByteSize = selfTotalLength * size_of_data_type;
    size_t indexByteSize = indexTotalLength * sizeof(DTYPE_INDEX);
    size_t srcByteSize = indexTotalLength * size_of_data_type;
    size_t tilingSize = sizeof(ScatterReduceTilingData);

    uint8_t *currentWorkspace = (uint8_t*)AscendC::GmAlloc(selfTotalLength * sizeof(DTYPE_INDEX)); // 通过框架获取workspace的指针，GetWorkspaceSizes入参为所需workspace的块数。当前限制使用一块。


#ifdef ASCENDC_CPU_DEBUG
    uint8_t *self = (uint8_t *)AscendC::GmAlloc(selfByteSize);
    uint8_t *index = (uint8_t *)AscendC::GmAlloc(indexByteSize);
    uint8_t *src = (uint8_t *)AscendC::GmAlloc(srcByteSize);

    ReadFile("./input/input_self.bin", selfByteSize, self, selfByteSize);
    ReadFile("./input/input_index.bin", indexByteSize, index, indexByteSize);
    ReadFile("./input/input_src.bin", srcByteSize, src, srcByteSize);

    ScatterReduceTilingData *tiling_data = (ScatterReduceTilingData *)AscendC::GmAlloc(tilingSize);
    tiling_data->blockNum = effectiveBlockNum;
    tiling_data->reduce = reduce;
    tiling_data->include_self = include_self;
    tiling_data->dim = dim;
    tiling_data->blockLength = blockLength;
    tiling_data->lastBlockLength = lastBlockLength;
    tiling_data->prevDimSize = prevDimSize;
    tiling_data->selfDimSize = selfDimSize;
    tiling_data->indexDimSize = indexDimSize;
    tiling_data->postDimSize = postDimSize;

    AscendC::SetKernelMode(KernelMode::AIV_MODE);

    ICPU_RUN_KF(scatter_reduce, effectiveBlockNum, self, index, src, currentWorkspace,
                *reinterpret_cast<ScatterReduceTilingData *>(tiling_data)); // use this macro for cpu debug

    WriteFile("./output/output_self.bin", self, selfByteSize);

    AscendC::GmFree((void *)self);
    AscendC::GmFree((void *)index);
    AscendC::GmFree((void *)src);
    AscendC::GmFree((void *)tiling_data);
#else
    CHECK_ACL(aclInit(nullptr));
    int32_t deviceId = 0;
    CHECK_ACL(aclrtSetDevice(deviceId));
    aclrtStream stream = nullptr;
    CHECK_ACL(aclrtCreateStream(&stream));

    ScatterReduceTilingData *tiling_data;
    uint8_t *selfHost, *indexHost, *srcHost;
    uint8_t *selfDevice, *indexDevice, *srcDevice;

    CHECK_ACL(aclrtMallocHost((void **)(&tiling_data), tilingSize));
    
    // Setup tiling data
    tiling_data->lastdim = lastdim;
    tiling_data->totalLength = totalLength;
    tiling_data->ALIGN_NUM = ALIGN_NUM;
    tiling_data->block_size = block_size;
    tiling_data->core_size = core_size;
    tiling_data->core_remain = core_remain;
    tiling_data->dim = dim;
    tiling_data->reduce = reduce;
    tiling_data->include_self = include_self;
    tiling_data->src_numel = src_numel;
    tiling_data->index_numel = index_numel;

    CHECK_ACL(aclrtMallocHost((void **)(&selfHost), selfByteSize));
    CHECK_ACL(aclrtMallocHost((void **)(&indexHost), indexByteSize));
    CHECK_ACL(aclrtMallocHost((void **)(&srcHost), srcByteSize));
    CHECK_ACL(aclrtMalloc((void **)&selfDevice, selfByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&indexDevice, indexByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&srcDevice, srcByteSize, ACL_MEM_MALLOC_HUGE_FIRST));

    ReadFile("./input/input_self.bin", selfByteSize, selfHost, selfByteSize);
    ReadFile("./input/input_index.bin", indexByteSize, indexHost, indexByteSize);
    ReadFile("./input/input_src.bin", srcByteSize, srcHost, srcByteSize);

    CHECK_ACL(aclrtMemcpy(selfDevice, selfByteSize, selfHost, selfByteSize, ACL_MEMCPY_HOST_TO_DEVICE));
    CHECK_ACL(aclrtMemcpy(indexDevice, indexByteSize, indexHost, indexByteSize, ACL_MEMCPY_HOST_TO_DEVICE));
    CHECK_ACL(aclrtMemcpy(srcDevice, srcByteSize, srcHost, srcByteSize, ACL_MEMCPY_HOST_TO_DEVICE));

    ACLRT_LAUNCH_KERNEL(scatter_reduce)(aivNum, stream, selfDevice, indexDevice, srcDevice, tiling_data);
    CHECK_ACL(aclrtSynchronizeStream(stream));

    CHECK_ACL(aclrtMemcpy(selfHost, selfByteSize, selfDevice, selfByteSize, ACL_MEMCPY_DEVICE_TO_HOST));
    WriteFile("./output/output_self.bin", selfHost, selfByteSize);

    CHECK_ACL(aclrtFree(selfDevice));
    CHECK_ACL(aclrtFree(indexDevice));
    CHECK_ACL(aclrtFree(srcDevice));
    CHECK_ACL(aclrtFreeHost(selfHost));
    CHECK_ACL(aclrtFreeHost(indexHost));
    CHECK_ACL(aclrtFreeHost(srcHost));
    CHECK_ACL(aclrtFreeHost(tiling_data));

    CHECK_ACL(aclrtDestroyStream(stream));
    CHECK_ACL(aclrtResetDevice(deviceId));
    CHECK_ACL(aclFinalize());
#endif
    return 0;
}